import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import os
import csv

class PupilAnnotator:
    def __init__(self, root):
        self.root = root
        self.root.title("Pupil Annotation Tool - CSV Version")

        # 初始化变量
        self.image_path = None
        self.image = None
        self.photo = None
        self.canvas = None
        self.circle = None
        self.radius = 10  # 初始半径
        self.center = None
        self.image_files = []  # 存储所有图片文件路径
        self.current_index = 0  # 当前图片的索引

        # 固定图片大小
        self.fixed_width = 800
        self.fixed_height = 600

        # 创建界面
        self.create_widgets()

        # 绑定事件
        self.root.bind("<MouseWheel>", self.on_mousewheel)
        self.root.bind("<Button-4>", self.on_mousewheel)
        self.root.bind("<Button-5>", self.on_mousewheel)
        self.root.bind("<Up>", lambda e: self.adjust_radius(1))
        self.root.bind("<Down>", lambda e: self.adjust_radius(-1))

    def adjust_radius(self, delta):
        self.radius = max(1, self.radius + delta)
        self.draw_circle()

    def create_widgets(self):
        # 创建画布框架，用于放置画布
        canvas_frame = tk.Frame(self.root)
        canvas_frame.pack(fill=tk.BOTH, expand=True)

        # 创建画布
        self.canvas = tk.Canvas(canvas_frame, cursor="cross", width=self.fixed_width, height=self.fixed_height)
        self.canvas.pack(fill=tk.BOTH, expand=True)
        self.canvas.bind("<Button-1>", self.on_click)

        # 创建按钮框架，用于放置按钮
        button_frame = tk.Frame(self.root)
        button_frame.pack(side=tk.BOTTOM, fill=tk.X)

        # 创建按钮
        btn_prev = tk.Button(button_frame, text="Previous Image", command=self.prev_image)
        btn_open = tk.Button(button_frame, text="Open Image", command=self.open_image)
        btn_save = tk.Button(button_frame, text="Save CSV", command=self.save_csv)
        btn_next = tk.Button(button_frame, text="Next Image", command=self.next_image)

        # 使用 pack 布局管理器将按钮均匀分布
        btn_prev.pack(side=tk.LEFT, fill=tk.X, expand=True)
        btn_open.pack(side=tk.LEFT, fill=tk.X, expand=True)
        btn_save.pack(side=tk.LEFT, fill=tk.X, expand=True)
        btn_next.pack(side=tk.LEFT, fill=tk.X, expand=True)

    def open_image(self):
        path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg *.jpeg *.png")])
        if not path:
            return

        # 获取当前目录下所有图片文件
        self.image_files = [os.path.join(os.path.dirname(path), f) for f in os.listdir(os.path.dirname(path)) if f.endswith(('.jpg', '.jpeg', '.png'))]

        # Normalize the paths
        normalized_path = os.path.normpath(path)
        normalized_image_files = [os.path.normpath(f) for f in self.image_files]

        # Find the index of the selected image
        self.current_index = normalized_image_files.index(normalized_path)

        self.image_path = path
        # 更新窗口标题
        self.root.title(f"Pupil Annotation Tool - {os.path.basename(path)}")
        original_image = Image.open(path)
        original_width, original_height = original_image.size
        # 调整图片大小以适应固定尺寸
        self.image = original_image.resize((self.fixed_width, self.fixed_height), Image.LANCZOS)
        self.photo = ImageTk.PhotoImage(self.image)

        # 清空画布
        self.canvas.delete("all")
        self.canvas.config(width=self.fixed_width, height=self.fixed_height)
        self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo)

        # 重置标注
        self.center = None
        self.radius = 10
        self.circle = None

        # 尝试加载标注
        self.load_annotation(original_width, original_height)

    def on_click(self, event):
        if not self.image:
            return

        self.center = (event.x, event.y)
        self.draw_circle()

    def on_mousewheel(self, event):
        if not self.center:
            return
        delta = event.delta if event.num != 4 else 1
        # 调整半径（Windows/Linux用event.delta，Mac用event.delta/0.1）
        delta = event.delta / 100 if delta > 10 else delta
        if event.state & 0x4:
            delta *= 3
        elif event.state & 0x1:
            delta = delta // abs(delta) if delta != 0 else 0
        self.radius = max(1, self.radius + int(delta))
        self.draw_circle()

    def draw_circle(self):
        if self.circle:
            self.canvas.delete(self.circle)
        if self.center:
            x, y = self.center
            self.circle = self.canvas.create_oval(
                x - self.radius, y - self.radius,
                x + self.radius, y + self.radius,
                outline="#FF4444", width=2,
                tags="circle"
            )

        self.canvas.delete("radius_text")
        if self.center:
            x, y = self.center
            self.canvas.create_text(
                x, y + self.radius + 15,
                text=f"R:{self.radius}px",
                fill="#44FF44",
                tags="radius_text"
            )

    def save_csv(self):
        if not self.image_path or not self.center:
            return

        # 生成CSV数据
        img_w, img_h = self.image.size
        data = {
            "filename": os.path.basename(self.image_path),
            "x_center": self.center[0],
            "y_center": self.center[1],
            "diameter": 2 * self.radius,
            "image_width": img_w,
            "image_height": img_h
        }

        # 保存CSV到同目录
        csv_path = os.path.splitext(self.image_path)[0] + ".csv"
        with open(csv_path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=data.keys())
            writer.writeheader()
            writer.writerow(data)

        print(f"CSV saved to {csv_path}")

    def prev_image(self):
        if self.image_files:
            self.current_index = (self.current_index - 1) % len(self.image_files)
            self.load_image()

    def next_image(self):
        if self.image_files:
            self.current_index = (self.current_index + 1) % len(self.image_files)
            self.load_image()

    def load_image(self):
        self.image_path = self.image_files[self.current_index]
        # 更新窗口标题
        self.root.title(f"Pupil Annotation Tool - {os.path.basename(self.image_path)}")
        original_image = Image.open(self.image_path)
        original_width, original_height = original_image.size
        # 调整图片大小以适应固定尺寸
        self.image = original_image.resize((self.fixed_width, self.fixed_height), Image.LANCZOS)
        self.photo = ImageTk.PhotoImage(self.image)

        # 清空画布
        self.canvas.delete("all")
        self.canvas.config(width=self.fixed_width, height=self.fixed_height)
        self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo)

        # 重置标注
        self.center = None
        self.radius = 10
        self.circle = None

        # 尝试加载标注
        self.load_annotation(original_width, original_height)

    def load_annotation(self, original_width, original_height):
        csv_path = os.path.splitext(self.image_path)[0] + ".csv"
        if os.path.exists(csv_path):
            with open(csv_path, 'r') as f:
                reader = csv.DictReader(f)
                # 读取所有行并获取最后一行（最新的标注）
                rows = list(reader)
                if rows:  # 确保有数据
                    row = rows[-1]  # 获取最后一行
                    # 计算缩放比例（使用原始图片尺寸而非调整后的尺寸）
                    width_ratio = self.fixed_width / float(row['image_width'])
                    height_ratio = self.fixed_height / float(row['image_height'])
                    # 缩放坐标
                    x_center = int(float(row['x_center']) * width_ratio)
                    y_center = int(float(row['y_center']) * height_ratio)
                    # 缩放半径（使用最小比例保持圆形）
                    radius = int(float(row['diameter']) / 2 * min(width_ratio, height_ratio))
                    self.center = (x_center, y_center)
                    self.radius = radius
                    self.draw_circle()


if __name__ == "__main__":
    root = tk.Tk()
    app = PupilAnnotator(root)
    root.mainloop()
